import requests
import json
from web_spider_bocha import search_web
from web_spider_deepseek import get_deepseek_response
from web_spider_kimi import get_kimi_response

import re
import jieba
import datetime

# Ollama API 地址
OLLAMA_API_URL = "http://localhost:11434"

# Deepseek模型
COMMON_MODEL_NAME = "deepseek-r1:1.5b"
AGRICULTURE_MODEL_NAME = "deepseek-r1-lvkong-v3"


# 系统提示
SYSTEM = """
你是一个严谨的AI助手，必须完全用中文回答问题，绝对不允许出现任何英文单词或短语。
你的回答必须详细，准确。
再次强调：所有回答必须完全使用中文，不得含有任何英文。
"""

# 整理搜索结果的系统提示
ORGANIZE_SYSTEM = """
你是一个专业的信息整理助手，你的任务是：
1. 整理和总结网络搜索返回的结果
2. 提取最相关和最有用的信息
3. 以简洁明了的方式组织信息
4. 保留关键的事实和数据
请完全使用中文回答，确保信息准确且易于理解，绝对不允许出现任何英文单词或短语。
"""

# 模板
def build_prompt(system_prompt, user_input, search_results=None):
    prompt = f"""<|system|>
{system_prompt}<|end|>
<|user|>
User: {user_input}<|end|>
<|assistant|>
Assistant:"""
    if search_results:
        prompt += f"\n以下是根据你的问题搜索到的实时信息：\n{search_results}\n"
    return prompt

# 使用deepseek-r1:1.5b整理搜索结果
def organize_search_results(query, search_results):
    """
    使用deepseek-r1:1.5b模型整理搜索结果
    
    :param query: 用户查询
    :param search_results: 搜索结果列表
    :return: 整理后的搜索结果文本，如果发生缓存溢出则返回None
    """
    print("正在整理搜索结果...")
    
    # 将搜索结果转换为文本
    # search_text = ""
    # for i, result in enumerate(search_results, 1):
    #     search_text += f"{i}. 标题: {result['name']}\n   链接: {result['url']}\n   摘要: {result['snippet']}\n\n"
    
    # kimi ai直接返回搜索结果
    search_text = search_results
    
    # 组织请求数据
    data = {
        "model": COMMON_MODEL_NAME,
        "prompt": f"""<|system|>
{ORGANIZE_SYSTEM}<|end|>
<|user|>
User: 请根据以下搜索结果，整理出与问题"{query}"相关的重要信息：

{search_text}<|end|>
<|assistant|>
Assistant:""",
        "stream": True
    }
    
    try:
        # 发送 POST 请求
        response = requests.post(f"{OLLAMA_API_URL}/api/generate", json=data, timeout=30, stream=True)
        
        # 处理流式响应
        if response.status_code == 200:
            full_response = ""
            skip_flag = False  # 跳过 <think> 标签
            for line in response.iter_lines():
                if line:
                    # 解析 JSON
                    chunk = line.decode("utf-8")
                    
                    try:
                        json_data = json.loads(chunk)
                        
                        # 检查是否存在缓存溢出错误信息
                        if "error" in json_data and any(err in json_data["error"].lower() for err in ["out of memory", "cache overflow", "memory limit", "buffer overflow"]):
                            print(f"缓存溢出错误：{json_data['error']}")
                            return None
                        
                        if json_data["response"] == "<think>" or json_data["response"] == "</think>":
                            skip_flag = not skip_flag
                            continue
                        
                        if "response" in json_data and not skip_flag:
                            full_response += json_data["response"]
                        
                        if json_data.get("done", False):
                            break
                    except json.JSONDecodeError as je:
                        print(f"JSON解析错误：{je}")
                        if "out of memory" in chunk.lower() or "cache overflow" in chunk.lower():
                            print("检测到缓存溢出错误")
                            return None
            
            return extract_chinese_from_string(full_response)
        else:
            print(f"请求失败，状态码：{response.status_code}")
            error_content = response.text
            if "out of memory" in error_content.lower() or "cache overflow" in error_content.lower():
                print("检测到缓存溢出错误")
                return None
            # 返回原始搜索结果
            return search_text
    except Exception as e:
        print(f"整理搜索结果时发生错误：{e}")
        if "out of memory" in str(e).lower() or "cache overflow" in str(e).lower():
            print("检测到缓存溢出错误")
            return None
        return search_text


# 用于提取汉字和标点符号
def extract_chinese_from_string(input_string):
    """
    提取字符串中的汉字、标点符号和数字。
    
    :param input_string: 原始字符串
    :return: 提取的汉字、标点符号和数字字符串
    """
    # 使用正则表达式提取所有汉字、中文标点符号和数字
    chinese_chars_and_punctuation = re.findall(r'[\u4e00-\u9fa5\u3000-\u303f\uff00-\uffef0-9]', input_string)
    # 将列表合并为字符串
    chinese_text = ''.join(chinese_chars_and_punctuation)
    # 返回汉字、标点符号和数字
    return chinese_text


# 提取关键词函数
def extract_keywords(text, top_n=10):
    """
    从文本中提取关键词
    
    :param text: 输入文本
    :param top_n: 返回前n个关键词
    :return: 关键词列表
    """
    # 使用jieba分词
    words = jieba.cut(text)
    # 过滤停用词和标点符号
    filtered_words = [word for word in words if len(word) > 1 and not re.match(r'[^\w\s]', word)]
    # 返回前top_n个词
    return filtered_words[:top_n]


# 处理时间信息函数
def process_time_info(query):
    """
    处理查询中的时间信息，将相对时间词转换为具体日期
    
    :param query: 用户查询
    :return: 处理后的查询和提取出的日期
    """
    today = datetime.datetime.now()
    yesterday = today - datetime.timedelta(days=1)
    tomorrow = today + datetime.timedelta(days=1)
    
    # 格式化日期
    today_str = today.strftime("%Y年%m月%d日")
    yesterday_str = yesterday.strftime("%Y年%m月%d日")
    tomorrow_str = tomorrow.strftime("%Y年%m月%d日")
    
    # 定义时间词映射
    time_mapping = {
        "今天": today_str,
        "今日": today_str,
        "现在": today_str,
        "当前日期": today_str,
        "昨天": yesterday_str,
        "昨日": yesterday_str,
        "明天": tomorrow_str,
        "明日": tomorrow_str
    }
    
    # 识别并替换相对时间词
    processed_query = query
    extracted_date = None
    
    for time_word, date_str in time_mapping.items():
        if time_word in query:
            processed_query = query.replace(time_word, date_str)
            extracted_date = date_str
            break
    
    return processed_query, extracted_date


# 主函数
def get_deepseek_response(user_input):
    if user_input == "退出" or user_input == "结束":
        print("感谢使用！")
        return
    
    # 步骤1: 获得用户输入
    refined_query = user_input # 字符串，用户输入的问题
    
    # 步骤2: 处理时间信息
    processed_query, extracted_date = process_time_info(refined_query)
    
    # 步骤3: 提取关键词
    keywords = extract_keywords(processed_query)
    print("提取的关键词：", keywords)
    
    # 如果有提取出的日期，确保它在关键词中
    if extracted_date and extracted_date not in keywords:
        keywords.append(extracted_date)
    
    # 步骤4: 判断是否需要联网搜索
    need_search = any(keyword in refined_query for keyword in ["最新", "实时", "今天", "现在", "当前", "天气", "预报", "新闻", "股市", "股票", "行情", "比赛", "比分", "赛事", "直播", "热点", "时事", "疫情", "突发", "事件", "政策", "法规", "汇率", "价格", "报价", "上市", "发布", "动态", "趋势", "展会", "会议", "大会", "节日", "活动", "展览", "发布会", "数据", "统计", "排名", "榜单", "排行"])
    search_results_text = None
    
    # 步骤5: 如果需要联网搜索，使用关键词进行搜索
    if need_search:
        print("正在联网搜索...")
        if keywords:
            search_keyword = " ".join(keywords)  # 使用提取的关键词
        else:
            search_keyword = processed_query  # 如果没有提取到关键词，使用处理后的查询
            
        # search_results = search_web(search_keyword)
        search_results = get_kimi_response(search_keyword)
        print("搜索结果：", search_results)
        
        if search_results and not isinstance(search_results, dict):  # 检查搜索结果是否有效
            # 步骤6: 使用deepseek-r1:1.5b整理搜索结果
            search_results_text = organize_search_results(processed_query, search_results)
            return search_results_text
    
    else:
        # 步骤6: 使用deepseek-r1-lvkong-v3生成最终回答
        # 组织请求数据
        data = {
            "model": AGRICULTURE_MODEL_NAME,
            "prompt": build_prompt(SYSTEM, refined_query, search_results_text),
            "stream": True
        }
        
        try:
            # 发送 POST 请求
            response = requests.post(f"{OLLAMA_API_URL}/api/generate", json=data, timeout=30, stream=True)
            
            # 处理流式响应
            if response.status_code == 200:
                full_response = ""
                skip_flag = False  # 跳过 <think> 标签
                for line in response.iter_lines():
                    if line:
                        # 解析 JSON
                        chunk = line.decode("utf-8")
                        
                        try:
                            json_data = json.loads(chunk)
                            
                            # 检查是否存在缓存溢出错误信息
                            if "error" in json_data and any(err in json_data["error"].lower() for err in ["out of memory", "cache overflow", "memory limit", "buffer overflow"]):
                                print(f"缓存溢出错误：{json_data['error']}")
                                return None
                            
                            if json_data["response"] == "<think>" or json_data["response"] == "</think>":
                                skip_flag = not skip_flag
                                continue
                            
                            if "response" in json_data and not skip_flag:
                                full_response += json_data["response"]
                            
                            if json_data.get("done", False):
                                break
                        except json.JSONDecodeError as je:
                            print(f"JSON解析错误：{je}")
                            if "out of memory" in chunk.lower() or "cache overflow" in chunk.lower():
                                print("检测到缓存溢出错误")
                                return None

                return full_response

            else:
                print(f"请求失败，状态码：{response.status_code}")
                error_content = response.text
                if "out of memory" in error_content.lower() or "cache overflow" in error_content.lower():
                    print("检测到缓存溢出错误")
                    return None
                print("API 返回：", response.text)
                return None
        
        except requests.exceptions.Timeout:
            print("请求超时，请检查网络或 API 服务是否正常。")
            return None

        except requests.exceptions.RequestException as e:
            print(f"请求发生错误：{e}")
            if "out of memory" in str(e).lower() or "cache overflow" in str(e).lower():
                print("检测到缓存溢出错误")
                return None
            return None


# if __name__ == '__main__':
#     while True:
#         input_text = input("请输入问题：")
#         if input_text == "":
#             break
#         result = get_deepseek_response(input_text)